# -*- coding: utf-8 -*-
# File : vul_scan_service.py
# Author: taoyahui
# Date : 2021/6/26
import time
from application_factory.extensions import scheduler
from tasks.vul_scan import vul_scan_integrity, call_vul_scan
from utils import mysql_util, common_util, job_util, log_util
import results
from application_factory.models import InventoryCls
import json


def vul_scan_entry_service(time_json, inventory_id):
    """
    漏扫功能入口服务
    :param date_time_str:
    :param inventory_id:
    :return:
    """
    time_dict = json.loads(time_json)
    exec_sql = f"select * from inventory where id = '{inventory_id}'"
    res = mysql_util.exec_sql_new(exec_sql)
    if res is None:
        return results.inventory_not_exist_result
    inventory = InventoryCls(res[0])
    if inventory.hostname is None:
        return results.get_hostname_error_result
    return vul_scan_service(inventory_ip=inventory.ip,
                            inventory_hostname=inventory.hostname,
                            time_dict=time_dict,
                            os=inventory.os)


# 启动漏洞扫描服务
# task add e.g
# scheduler.add_job(func=execute_strategy,
#                       id=f"{add_strategy.strategy_id}",
#                       kwargs={'strategy_id': add_strategy.strategy_id, 'is_test': is_test},
#                       trigger=trigger_type, **cron_rel,
#                       replace_existing=True)
def vul_scan_service(inventory_ip, inventory_hostname, os, time_dict={'type':'每天', 'time':'2021-07-11T02:00:00.000Z'}):
    tm_hour = 2
    tm_min = 0
    cron_rel = {}
    time_str = time_dict['time']
    time_parse = common_util.str_to_time_format(time_str)
    tm_hour = time_parse.hour
    tm_min = time_parse.minute
    if time_dict['type'] == '每天':
        cron_rel = dict(minute=f"{tm_min}", hour=f"{tm_hour}")
        pass
    elif time_dict['type'] == '每周':
        cron_rel = dict(minute=f"{tm_min}",
                        hour=f"{tm_hour}",
                        day_of_week=common_util.week_chinese_to_int(time_dict['week']))
    elif time_dict['type'] == '每月':
        cron_rel = dict(minute=f"{tm_min}",
                        hour=f"{tm_hour}",
                        month=time_dict['month'])

    # if date_time_str is not None:
    #     scan_time = time.strptime(date_time_str, "%H:%M")
    #     # 扫描的小时
    #     tm_hour = scan_time.tm_hour
    #     # 扫描的分钟
    #     tm_min = scan_time.tm_min
    job = trigger_vul_scan_job(hostname=inventory_hostname, ip=inventory_ip, os=os, cron_rel=cron_rel)
    return results.success_result


# 漏洞扫描任务
def trigger_vul_scan_job(hostname, ip, cron_rel, os, trigger_type='cron'):
    """
    启动漏洞扫描任务
    :param hostname:
    :param ip:
    :param cron_rel:
    :param trigger_type:
    :return:
    """
    if 'windows' not in os.lower() or '未知' not in os:
        log_util.logger.info(f"trigger_vul_scan_job -> Start to vul scan os is {os}")
        job = scheduler.add_job(func=call_vul_scan,
                                id=f"Vul-Scan-{hostname}",
                                kwargs={'ip': ip, 'hostname': hostname},
                                trigger=trigger_type, **cron_rel,
                                replace_existing=True)
        job_util.add_scan_job(f"Vul-Scan-{hostname}", hostname, ip, cron_rel)
    return job


def remove_safe_detect_service(inventory_id):
    """
    删除资产的漏扫功能
    :param inventory_id:
    :return:
    """
    exec_sql = f"select * from inventory where id = '{inventory_id}'"
    res = mysql_util.exec_sql_new(exec_sql)
    if res is None:
        return results.inventory_not_exist_result
    inventory = InventoryCls(res[0])
    res = job_util.remove_scan_job(inventory.hostname)
    if res['code'] != 200:
        return res
    scheduler.remove_job(f"Vul-Scan-{inventory.hostname}")

    return results.success_result


# 立即执行漏洞扫描的函数
def vul_scan_immediately(inventory_id):
    exec_sql = f"select * from inventory where id = {inventory_id}"
    res = mysql_util.exec_sql_new(exec_sql)
    if res is None:
        return results.inventory_not_exist_result
    inventory = InventoryCls(res[0])
    if 'windows' not in inventory.os.lower():
        job = scheduler.add_job(func=call_vul_scan,
                                id=f"Vul-Scan-{inventory.hostname}-{int(time.time())}",
                                kwargs={'ip': inventory.ip, 'hostname': inventory.hostname}
                                )
    return results.success_result


if __name__ == '__main__':
    # vul_scan_service("22:30")
    pass
